'''
Created on Apr 12, 2011

@author: rreyes
'''

from yacf.MiddleEnd.Loop.Analysis import ParametrizeLoopTool
from yacf.Backends.Common.Mutators.AbstractMutator import AbstractMutator

from yacf.Tools.Tree import ReplaceTool, getContainerAttribute
from yacf.Frontend.Shortcuts import getCurrentLanguageAst
from yacf.Backends.Common.Visitors.GenericVisitors import GenericFilterVisitor
c_ast = getCurrentLanguageAst()

from yacf.Frontend.Omp import omp_ast

class LoopUnrollFilter(GenericFilterVisitor):
    """ Returns the first node with the given attribute
    """
    def __init__(self):
        def condition(node):
            if type(node) == omp_ast.llcUnroll:
                return True
            return False
        super(LoopUnrollFilter, self).__init__(condition_func = condition)

class LoopUnrollMutator(AbstractMutator):
    def __init__(self, factor):
        self._factor = factor
        super(LoopUnrollMutator, self).__init__()

    
    def filter(self, ast):
        """  """
        raise NotImplemented

    def filter_iterator(self, ast):
        """ Fast filter  """
        return NotImplemented

    def fast_filter(self, ast):
        """ Fast filter , looking for binary expressions """
        return LoopUnrollFilter().dfs_iter(ast)
    
    def mutatorFunction(self, loopUnroll_directive_ast):
        
        for_ast = loopUnroll_directive_ast.loop
        
        assert isinstance(for_ast, c_ast.For)
        
        # Parametrize
        parametrizer = ParametrizeLoopTool()
        loop_parameters = parametrizer.apply(for_ast,)
        
        import copy
        # Note the deepcopy operation instead a copy or a bare assignment
        old_parameters = copy.deepcopy(loop_parameters)
        # Modify parameters with the unrolling
        # 1. Replace last_iteration by last_iteration/factor
        
        remainder_ast = c_ast.BinaryOp(op = '%',
                                       left = loop_parameters['last_iteration'],
                                       right = c_ast.Constant(type = 'int', value = str(self._factor))
                                      )
        
        compute_unroll_factor_ast = c_ast.BinaryOp(op = '-',
                                                   left =  loop_parameters['last_iteration'],
                                                   right = remainder_ast
                                                   )
        
        loop_parameters['last_iteration'] = compute_unroll_factor_ast
        

        
        # 2. Replace all occurrences of last_iteration by compute_last_iteration within the condition node
        max_in_expr_ast = parametrizer.symbolic_max(loop_parameters['cond_node'])
        # NOTE: Using setattr to REPLACE from the NODE . Using an assignment (=) will create a new instance
        setattr(max_in_expr_ast.parent, getContainerAttribute(max_in_expr_ast), compute_unroll_factor_ast)
        
        
        inside_unroll_compound_ast = c_ast.Compound(block_items = [])
        
        # Add a new stmt for each iteration that we want to unroll
        i = (self._factor - 1)
        while i >= 0:
            inside_unroll_compound_ast.block_items.append(for_ast.stmt.__copy__())
            if i > 0:
                inside_unroll_compound_ast.block_items.append(loop_parameters['iterator'])
            i -= 1
        
        
        # Update parameters

        parametrizer.update(loop_parameters)
        
        # print  "===> " + str(parametrizer) + "<===="
        
        # Build a loop from the parameter
        new_for_ast = parametrizer.buildLoop()        
        setattr(new_for_ast, 'stmt', inside_unroll_compound_ast)
        
        
        # Build a forLoop for the remaining parameters 
        #   Note: A foorLoop is better if we want to kernelize later       
        old_parameters['init_node'] = c_ast.Assignment(op = '=',
                                                       lvalue = old_parameters['loop_variable'],
                                                       rvalue = old_parameters['loop_variable'],
                                                       )
       
        parametrizer.update(old_parameters)
        remaining_loop_ast = parametrizer.buildLoop()
        setattr(remaining_loop_ast, 'stmt', for_ast.stmt.__copy__() )
        
        unroll_for_block_ast = c_ast.Compound(block_items = [])
        unroll_for_block_ast.block_items.append(new_for_ast)
        unroll_for_block_ast.block_items.append(remaining_loop_ast)
        # Replace the forLoop with the unrolled compound statement
        ReplaceTool(new_node = unroll_for_block_ast, 
                    old_node = for_ast.parent.parent
                    ).apply(for_ast.parent.parent.parent, 
                            'block_items'
                            )
        
        # print str(parametrizer)
        # New node
        # print str(new_for_ast)

